import asyncio
import websockets
from fastapi import FastAPI, WebSocket
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
import uvicorn
import json
from typing import List, Set
import logging
from fastapi import WebSocket, WebSocketDisconnect
from websockets.exceptions import ConnectionClosedError
import random
import time

app = FastAPI()

# 静态文件服务
app.mount("/static", StaticFiles(directory="static"), name="static")

# WebSocket连接管理器
class ConnectionManager:
    def __init__(self):
        self.active_connections: Set[WebSocket] = set()
        self._lock = asyncio.Lock()
        
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        async with self._lock:
            self.active_connections.add(websocket)
        logging.info(f"New connection from {websocket.client}")
        
    async def disconnect(self, websocket: WebSocket):
        async with self._lock:
            if websocket in self.active_connections:
                self.active_connections.remove(websocket)
                logging.info(f"Connection closed: {websocket.client}")
            else:
                logging.warning(f"Attempted to remove non-existent connection")
                
    async def broadcast(self, data: dict):
        dead_connections = []
        async with self._lock:
            connections = list(self.active_connections)
            
        for connection in connections:
            try:
                await connection.send_json(data)
            except (WebSocketDisconnect, ConnectionClosedError, RuntimeError) as e:
                logging.warning(f"Broadcast failed: {type(e).__name__}")
                dead_connections.append(connection)
            except Exception as e:
                logging.error(f"Unexpected broadcast error: {e}")
                dead_connections.append(connection)
                
        if dead_connections:
            async with self._lock:
                self.active_connections.difference_update(dead_connections)

manager = ConnectionManager()

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 模拟数据生成任务 - 调整为20毫秒间隔
async def data_generator():
    try:
        while True:
            try:
                # 生成模拟数据
                data = {
                    "type": "waveform",
                    "data": [random.randint(-32768, 32767) for _ in range(10)],
                    "timestamp": time.time(),
                    "sample_rate": 1000,  # 1000Hz采样率
                    "interval_ms": 20     # 20毫秒间隔
                }
                await manager.broadcast(data)
                await asyncio.sleep(0.02)  # 20毫秒间隔
                
            except asyncio.CancelledError:
                logger.info("Data generator task cancelled")
                raise
            except Exception as e:
                logger.error(f"Data generation error: {e}")
                await asyncio.sleep(1)  # 错误时暂停1秒
                
    finally:
        logger.info("Data generator stopped")

# WebSocket端点 - 优化心跳处理
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            try:
                # 接收消息，设置合理的超时时间
                data = await asyncio.wait_for(
                    websocket.receive_text(),
                    timeout=30.0  # 30秒超时
                )
                
                if data.lower() == "ping":
                    await websocket.send_text("pong")
                    logger.debug("Received ping, sent pong")
                elif data.lower() == "pong":
                    logger.debug("Received pong from client")
                    continue
                else:
                    # 处理其他消息
                    logger.debug(f"Received message: {data}")
                    
            except asyncio.TimeoutError:
                # 发送ping检查连接状态
                try:
                    await websocket.send_text("ping")
                    # 等待pong响应
                    response = await asyncio.wait_for(
                        websocket.receive_text(),
                        timeout=3.0
                    )
                    if response.lower() != "pong":
                        logger.warning("Invalid heartbeat response")
                        break
                    else:
                        logger.debug("Heartbeat successful")
                except asyncio.TimeoutError:
                    logger.warning("Heartbeat timeout - client not responding")
                    break
                except (WebSocketDisconnect, ConnectionClosedError):
                    break
                    
    except WebSocketDisconnect as e:
        logger.info(f"Client cleanly disconnected: code={e.code}")
    except ConnectionClosedError as e:
        logger.info(f"Client connection closed: code={e.code}")
    except Exception as e:
        logger.error(f"Unexpected WebSocket error: {type(e).__name__}: {e}")
    finally:
        await manager.disconnect(websocket)

@app.get("/connections")
async def get_connections():
    return {"active_connections": len(manager.active_connections)}

@app.get("/status")
async def get_status():
    return {
        "status": "running",
        "data_interval_ms": 20,
        "sample_rate_hz": 1000,
        "active_connections": len(manager.active_connections)
    }

# 提供HTML页面
@app.get("/")
async def get():
    with open("static/index.html") as f:
        return HTMLResponse(f.read())

# 添加定期连接检查
async def connection_health_check():
    while True:
        await asyncio.sleep(30)
        async with manager._lock:
            current_connections = list(manager.active_connections)
        
        if current_connections:
            logger.info(f"Health check: {len(current_connections)} active connections")
        
        for connection in current_connections:
            try:
                # 轻量级心跳检查
                await connection.send_json({
                    "type": "heartbeat", 
                    "timestamp": time.time()
                })
            except (WebSocketDisconnect, ConnectionClosedError, RuntimeError):
                logger.debug("Health check failed for one connection")

async def main():
    # 启动数据生成任务
    asyncio.create_task(data_generator())
    # 启动健康检查
    asyncio.create_task(connection_health_check())
    
    # 启动Web服务器
    config = uvicorn.Config(
        app, 
        host="0.0.0.0", 
        port=8000,
        log_level="info"
    )
    server = uvicorn.Server(config)
    await server.serve()
    
if __name__ == "__main__":
    asyncio.run(main())